
from opencv_usage.av_base import AVBase
import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from skimage import io

class AVOperation(AVBase):
    def __init__(self, isShow=False):
        self.show = isShow

    def rotateIt(self, infile, outfile):
        print("do rotate")
        indata = cv2.imread(infile)
        outmat = cv2.rotate(indata, cv2.ROTATE_180)
        if self.show:
            self.testShowTwo(indata, outmat)
        else:
            cv2.imwrite(outfile, outmat)
        return

    def rorateImageData(self, infile):
        indata = cv2.imread(infile)
        result = cv2.rotate(indata, cv2.ROTATE_180)
        return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)

    def testShowTwo(self, imgleft, imgright):
        fig, axs = plt.subplots(1,2)
        imgleft = cv2.cvtColor(imgleft, cv2.COLOR_RGB2BGR)
        imgright = cv2.cvtColor(imgright, cv2.COLOR_RGB2BGR)
        axs[0].imshow(imgleft)
        axs[0].set_title('Left')
        axs[1].imshow(imgright)
        axs[1].set_title('Right')
        plt.tight_layout()
        plt.show()

    def testEqualizeHist(self, infile, outfile):
        print("do equalize hist")
        indata = cv2.imread(infile, cv2.IMREAD_ANYCOLOR)
        (b, g, r) = cv2.split(indata)
        bH = cv2.equalizeHist(b)
        gH = cv2.equalizeHist(g)
        rH = cv2.equalizeHist(r)
        result = cv2.merge((bH, gH, rH))
        cv2.imwrite(outfile, result)
        if self.show:
            self.testShowTwo(indata, result)

    # 写一个函数，输入一个图片文件 ，计算其直方图规定化，输出到另一个图片地址中，同时在 subplot 国
    # 同时显示原图和处理后的图片
    def testHistSpec(self, infile, outfile):
        indata = cv2.imread(infile, cv2.IMREAD_ANYCOLOR)
        outdata = cv2.imread(outfile, cv2.IMREAD_ANYCOLOR)
        result = self.testHistSpecImpl(indata, outdata)
        self.testShowTwo(indata, result)

    def testHistSpecData(self, infile, outfile):
        indata = cv2.imread(infile, cv2.IMREAD_ANYCOLOR)
        outdata = cv2.imread(outfile, cv2.IMREAD_ANYCOLOR)
        result = self.testHistSpecImpl(indata, outdata)
        return cv2.cvtColor(result, cv2.COLOR_BGR2RGB)

    def testHistSpecColor(self, infile, outfile):
        indata = cv2.imread(infile)
        outdata = cv2.imread(outfile)

        b_i, g_i, r_i = cv2.split(indata)
        b_o, g_o, r_o = cv2.split(outdata)
        r_i = self.testHistSpecImpl(b_i, b_o)
        g_i = self.testHistSpecImpl(g_i, g_o)
        b_i = self.testHistSpecImpl(r_i, r_o)
        result = cv2.merge([r_i, g_i, b_i])
        self.testShowTwo(indata, result)

    def testHistSpecImpl(self, indata, outdata):
        print("do hist sure")
        source_hist, _ = np.histogram(indata.ravel(), bins=256, range=[0, 256])
        reference_hist, _ = np.histogram(outdata.ravel(), bins=256, range=[0, 256])

        # 计算累积分布函数
        source_cdf = source_hist.cumsum()
        reference_cdf = reference_hist.cumsum()

        # 避免除以零
        reference_cdf_m = reference_cdf.max()
        if reference_cdf_m != 0:
            reference_cdf = reference_cdf / float(reference_cdf_m)
        else:
            reference_cdf = reference_cdf

        source_cdf_m = source_cdf.max()
        if source_cdf_m != 0:
            source_cdf = source_cdf / float(source_cdf_m)
        else:
            source_cdf = source_cdf

        # 创建直方图映射
        s_z = np.interp(source_cdf, reference_cdf, range(256))
        s_z = np.round(s_z).astype(np.uint8)

        # 应用映射
        matched = s_z[indata.ravel()]
        result = matched.reshape(indata.shape)
        return result

    def testImageCutData(self, infile):
        indata = cv2.imread(infile, cv2.IMREAD_ANYCOLOR)
        T = np.array([[1, 0, 1],[0.8, 1, 0]])
        sheard_image = cv2.warpAffine(indata, T, indata.shape[:2])
        return indata,sheard_image

    def testImageCutShow(self, infile):
        indata, result = self.testImageCutData(infile)
        self.testShowTwo(indata, result)

    def testImagePerspectiveShow(self, infile):
        image = cv2.imread(infile)
        # 定义源点和目标点
        # 假设源点是图像的四个角点
        src_points = np.float32([[0, 0], [image.shape[1]-1, 0], [0, image.shape[0]-1], [image.shape[1]-1, image.shape[0]-1]])
        # 定义目标点，这里我们只是将图像向右移动
        dst_points = np.float32([[100, 200], [image.shape[1]-1, 0], [0, image.shape[0]-1], [image.shape[1]-1, image.shape[0]-1]])

        # 计算透视变换矩阵
        M = cv2.getPerspectiveTransform(src_points, dst_points)

        # 应用透视变换
        warped_image = cv2.warpPerspective(image, M, (image.shape[1]+200, image.shape[0]))
        self.testShowTwo(image, warped_image)

    def testImageDataToFile(self, infile, outfile):
        img = io.imread(infile)
        df = pd.DataFrame(img.flatten())
        df.to_excel(outfile, index = False)

    def testImageGammaCorrection(self, infile):
        img = cv2.imread(infile)

        gamma = 2.6
        inv_gamma = 1.0 / gamma
        table = (np.arange(256.0/255, step=1.0/255) ** gamma) * 255
        table = np.uint8(table)

        gamma_corrected = cv2.LUT(img, table)

        self.testShowTwo(img, gamma_corrected)

    def testImageMerge(self, infile, outfile):
        img_in = cv2.imread(infile)
        img_out = cv2.imread(outfile)
        final_image = cv2.addWeighted(img_in, 0.7, img_out, 0.3, 0.0)
        self.testShowTwo(img_in, final_image)

    def testImageEnhance(self, infile, outfile):
        img_in = cv2.imread(infile)
        img_float = img_in.astype(np.float32)
        adjusted_img = cv2.add(img_float, 20)
        adjusted_img = cv2.multiply(adjusted_img, 0.5)
        adjusted_img = np.clip(adjusted_img, 0, 255).astype(np.uint8)
        if self.show:
            self.testShowTwo(img_in, adjusted_img)
        else:
            cv2.imwrite(outfile, adjusted_img)

    def testConvolution(self, infile, outfile):
        img_in = cv2.imread(infile, cv2.IMREAD_GRAYSCALE)
        kernel = np.array([[1, 2, 1], [2, 4, 2], [1, 2, 1]])

        # 进行卷积操作
        result = cv2.filter2D(img_in, -1, kernel)
        self.testShowTwo(img_in, result)
